"""
2.	傅里叶变换告诉我们图像的时域卷积等价于对应的频域相乘，本题在时域和频域同时研究图像的傅里叶变换。按照下列要求，完成相应内容（共30分）
"""

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

spr = 2
spc = 3
spn = 0
plt.figure(figsize=(12, 6))


def my_imshow(title, img_data, color='gray'):
    """Utility function to show images."""
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.axis('off')
    if 'gray' == color:
        plt.imshow(img_data, cmap='gray')
    else:
        img_data = cv.cvtColor(img_data, cv.COLOR_BGR2RGB)
        plt.imshow(img_data)


def my_debug(arr, name):
    """Utility function for debugging."""
    print(f'{name}: {arr.dtype}, {arr.shape}, min = {arr.min()}, max = {arr.max()}, mean = {arr.mean()}')


# ①	读取图像文件（football.jpg）标记为img，并显示图片
path = 'data/images/football.jpg'
img = cv.imread(path, cv.IMREAD_GRAYSCALE)
print(f'打印图片维度和尺度: {img.shape}')
H, W = img.shape
H2 = H // 2
W2 = W // 2
my_imshow('original', img)
img_ = img.copy()

# ②	对img进行离散傅里叶变换，计算并显示幅度谱（magnitude_spectrum）
f = np.fft.fft2(img)  # fourier
fshift = np.fft.fftshift(f)  # fourier shifted
fshift_ = fshift.copy()
EPS = 1e-10  # Epsilon added to np.log's parameter
magnitude_spectrum = np.array(np.log(abs(fshift) + EPS))
my_debug(magnitude_spectrum, 'magnitude_spectrum')
my_imshow('magnitude_spectrum', magnitude_spectrum)

# ③	在时域确定均值滤波核（3X3）,计算img对该核的卷积conv_mean并显示
kernel_mean = np.ones((3, 3), dtype=np.float32)
conv_mean = cv.filter2D(img, cv.CV_32F, kernel_mean)
cv.normalize(conv_mean, conv_mean, 0., 255., cv.NORM_MINMAX)
conv_mean = conv_mean.astype(np.uint8)
my_imshow('conv_mean', conv_mean)

# ④	在时域确定拉普拉斯核,计算img对该核的卷积conv_laplacian并显示
kernel_lap = np.float32([
    [0, 1, 0],
    [1, -4, 1],
    [0, 1, 0],
])
conv_laplacian = cv.filter2D(img, cv.CV_32F, kernel_lap)
cv.normalize(conv_laplacian, conv_laplacian, 0., 255., cv.NORM_MINMAX)
conv_laplacian = conv_laplacian.astype(np.uint8)
my_imshow('conv_laplacian', conv_laplacian)

# ⑤	在频域，幅度谱（magnitude_spectrum）通过低通滤波器，计算傅里叶反变换img_back1并显示结果
fshift4lpf = fshift_.copy()  # fourier shifted for LPF
mask = np.zeros((H, W), dtype=np.float32)
HALF_SCALE = 30
mask[H2 - HALF_SCALE:H2 + HALF_SCALE, W2 - HALF_SCALE:W2 + HALF_SCALE] = 1.
fshift4lpf *= mask
f4lpf = np.fft.ifftshift(fshift4lpf)  # fourier for LPF
img_back1 = np.fft.ifft2(f4lpf)
img_back1 = abs(img_back1)
my_imshow('img_back1', img_back1)

# ⑥	在频域，幅度谱（magnitude_spectrum）通过高通滤波器，计算傅里叶反变换img_back2并显示结果
fshift4hpf = fshift_.copy()  # fourier shifted for hpf
HALF_SCALE = 30
fshift[H2 - HALF_SCALE:H2 + HALF_SCALE, W2 - HALF_SCALE:W2 + HALF_SCALE] = 0
f4hpf = np.fft.ifftshift(fshift)  # fourier for hpf
img_back2 = np.fft.ifft2(f4hpf)
img_back2 = abs(img_back2)
my_imshow('img_back2', img_back2)

# show all plotting
plt.show()
